/**
 * @file add_custom.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "kernel_operator.h"
#include "softmaxv2_tiling.h"
#include "softmaxv2.h"

using namespace AscendC;

__aicore__ inline void CopyTiling(SoftmaxV2TilingData* tiling, GM_ADDR tilingGM) {
    uint32_t* ptr = reinterpret_cast<uint32_t*>(tiling);
    auto tiling32 = reinterpret_cast<__gm__ uint32_t*>(tilingGM);
  
    for (int i = 0; i < sizeof(SoftmaxV2TilingData) / sizeof(uint32_t); i++, ptr++) {
      *ptr = *(tiling32 + i);
    }
    return;
}

extern "C" __global__ __aicore__ void softmaxv2_custom(GM_ADDR srcGM, GM_ADDR dstGM, GM_ADDR softmaxMaxGM, 
                                                GM_ADDR softmaxSumGM, GM_ADDR tilingGM)
{
    TPipe pipe;
    SoftmaxV2TilingData softmaxV2Tiling;
    CopyTiling(&softmaxV2Tiling, tilingGM);

    SoftmaxV2<float,float> op;
    op.Init(srcGM, dstGM, softmaxMaxGM, softmaxSumGM, softmaxV2Tiling, &pipe);
    op.Process();
}
